;; Protocol for multiarrays

;; by Konrad Hinsen
;; last updated April 15, 2010

;; Copyright (c) Konrad Hinsen, 2009-2010. All rights reserved.  The use
;; and distribution terms for this software are covered by the Eclipse
;; Public License 1.0 (http://opensource.org/licenses/eclipse-1.0.php)
;; which can be found in the file epl-v10.html at the root of this
;; distribution.  By using this software in any fashion, you are
;; agreeing to be bound by the terms of this license.  You must not
;; remove this notice, or any other, from this software.

(ns multiarray.protocol
  (:use [multiarray.utility]))

(set! *warn-on-reflection* true)

;; Array constructors

; Constructors cannot be part of the array protocol because they
; take no array arguments. They are summarized here for completeness,
; but they are defined in each implementation module.

; Zeros returns an array of a given shape with all elements set to zero.
; An optional second argument specifies the type of the array elements,
; the default being Integer/TYPE.
; Zeros is defined as a first constructor mainly because it is so simple
; to implement. It may go away later if it turns out not to be important
; in practice.
(comment
  (zeros shape-vector)
  (zeros shape-vector element-type))

; ma-range works like range, but produces a multiarray rather than a seq
(comment
  (ma-range end)
  (ma-range start end)
  (ma-range start end step))

; into-ma converts anything that implements the MultiArray protocol into
; a multiarray
(comment
  (into-ma object))

;;
;; The multiarray protocol
;;

(defprotocol MultiArrayProtocol
  "Multidimensional array protocol"
  (shape [x]
  "returns a vector containing the number of elements along
   each dimension of x")
  (rank [x]
  "returns the number of dimensions of x")
  (elements [x]
  "returns a seq of all array elements")
  (reshape [x shape]
  "returns a multiarray of the given shape having the same elments as x")
  (sample [x #^Integer dim #^Integer start #^Integer end #^Integer step]
  "samples the array along dimension dim from index start to index end
   incrementing by step")
  (slice [x #^Integer dim #^Integer index]
  "selects a slice at index along dimension dim")
  (transpose [x] [x dim-permutation]
  "transpose the array such that the dimensions are rearranged
   as specified by dim-permutation. The default is the reverse of
   the current dimensions."))

;;
;; The low-level multiarray protocol for efficient data sharing between
;; different implementations.
;;

(defprotocol MultiArrayLowLevelProtocol
  "Multidimensional array protocol"
  (element-class [x]
  "returns the class of the elements of x")
  (element-array [x]
  "returns a Java array containing all array elements"))

;;
;; Functions based on the multiarray protocol. They may eventually be
;; integrated into the protocol, or into a separate protocol, in order to
;; permit optimization for specific implementations.
;;

(defn view [x & indices]
  "Return a view on a subset of x defined by the sequence of indices.
   The items of indices can be integers, specifying a slice operation,
   or one- to three-element vectors, specifying a sample operation."
  (letfn [(sample-args
	   ([end]            [0 end 1])
	   ([start end]      [start end 1])
	   ([start end step] [start end step]))
	  (view1 [[a dim] index]
		 (assert (< dim (rank a)))
		 (cond (vector? index)
  		         (let [[start end step] (apply sample-args index)]
			   [(sample a dim start end step) (inc dim)])
		       (map? index)
		         (let [{start :from end :to step :by
				:or {start 0 step 1 end nil}} index
				end  (if (nil? end)
				       (get (shape a) dim)
				       end)]
			   [(sample a dim start end step) (inc dim)])
		       (integer? index)
		         [(slice a dim index) dim]
		       (= :all index)
		         [a (inc dim)]
		       (= :reverse index)
		         [(sample a dim (dec (get (shape a) dim)) -1 -1)
			  (inc dim)]
		       (= :new index)
		         [(reshape a (vec (insert (shape a) dim 1))) dim]
		       :else
		         (throw (Exception.
				 (str "Invalid index expression" index)))))]
    (first (reduce view1 [x 0] indices))))

;;
;; Maps containing useful default implementations for some protocol functions
;;

(defn- not-implemented
  [& args]
  (throw (Exception. "not implemented")))

(def multiarray-method-defaults
     {:shape         not-implemented
      :rank          (fn [x] (count (shape x)))
      :elements      not-implemented
      :reshape       not-implemented
      :sample        not-implemented
      :slice         not-implemented
      :transpose     not-implemented})

(def multiarray-lowlevel-method-defaults
     {:element-class not-implemented
      :element-array (fn [x]
		       (let [el-class (element-class x)
			     elements (elements x)]
			 (into-array el-class elements)))})

;;
;; The implementation for all non-array objects. They are considered
;; rank-0 arrays containing themselves.
;;

(extend Object
  MultiArrayProtocol
  (merge multiarray-method-defaults
	 {:shape         (fn [x] [])
	  :rank          (fn [x] 0)
	  :elements      list
	  :transpose     identity})
  MultiArrayLowLevelProtocol
  (merge multiarray-lowlevel-method-defaults
	 {:element-class (fn [x] (common-element-class (list x)))}))
 

(set! *warn-on-reflection* false)
